#!/usr/bin/env python
# coding: utf-8

# In[1]:


# 针对课上的代码，重写为ipynb文件，方便大家交互查看
import sys
import cv2
from random import randint


# In[2]:


# 跟踪器的选择列表
trackerTypes = ['BOOSTING', 'MIL', 'KCF','TLD', 'MEDIANFLOW', 'GOTURN', 'MOSSE', 'CSRT']

def createTrackerByName(trackerName):
    if trackerName in trackerTypes:
        tracker = eval(f'cv2.Tracker{trackerName}_create()')
    else:
        tracker = None
        print(f'Incorrect tracker name!\nAvailable trackers are:\n{[t for t in trackerTypes]}')
    return tracker# 就是这行在课上没加，导致没有返回跟踪器，程序直接跳出

# OPENCV 最大的缺点就是基本没有报错信息，或者报错信息很难看懂，这说明作者团队异常处理的不是很好

# 载入视频
videoPath = "siderwalk.mp4"

cap = cv2.VideoCapture(videoPath)

# 读取首帧
success, frame = cap.read()
if not success:
    print('Failed to read video')

# 读取首帧中的检测框
bboxes = []
colors = [] 

# 回车确认单次框，按q确认全部完成，进入跟踪阶段
while True:

    bbox = cv2.selectROI('MultiTracker', frame)
    bboxes.append(bbox)
    colors.append((randint(0, 255), randint(0, 255), randint(0, 255)))
    print("Press q to quit selecting boxes and start tracking")
    print("Press any other key to select next object")
    k = cv2.waitKey(0) & 0xFF
    if k == 113:  # ord('q')==113
        break

# 选择tracker类型
trackerType = "KCF"   

# 创建multi tracker
multiTracker = cv2.MultiTracker_create()


for bbox in bboxes:
    multiTracker.add(createTrackerByName(trackerType), frame, bbox)

# 开始处理视频
while cap.isOpened():
    success, frame = cap.read()
    if not success:
        break
  
 # 更新每一帧的跟踪结果
    success, boxes = multiTracker.update(frame)

    # 绘制矩形框
    for i, newbox in enumerate(boxes):
        p1 = (int(newbox[0]), int(newbox[1]))
        p2 = (int(newbox[0] + newbox[2]), int(newbox[1] + newbox[3]))
        cv2.rectangle(frame, p1, p2, [255,0,255], 2, 1)

    # 刷新显示结果
    cv2.imshow('MultiTracker', frame)


    # 中途退出Esc
    if cv2.waitKey(1) & 0xFF == 27:  # Esc pressed
        break


# In[ ]:




